######################
# To get gradient matrix (Empirical Estimating Functions) from mnlogit object
#The estimating function (or score function) for a model 
#is the derivative of the objective function with respect 
#to the parameter vector. The empirical estimating functions 
#is the evaluation of the estimating function at the observed data 
#(n observations) and the estimated parameters (of dimension k)
#########################
estfunjg <- function(mnlogitobject,fm) {
  #source(paste(codes,"parseFormulajg.R",sep=""))
  formula <- parseFormulajg(fm) 
  choiceVar <- "_Alt_Indx_"
  data <- mnlogitobject$model #ml.m.w$model
  data[[choiceVar]] <- attr(data, "index")$alt # attach to 'data'
  data <- data[order(data[[choiceVar]]), ]
  choice.set <- unique(data[[choiceVar]])
  K <- length(choice.set) # number of choices
  N <- nrow(data)/K
  # Obtain response vector as a vector of 0,1
  respVec <- data[[attr(formula, "response")]]
  if (is.factor(respVec)) respVec <- droplevels(respVec)
  respVec <- as.numeric(respVec)
  min.respVec <- min(respVec)
  spread <- max(respVec) - min.respVec
  if (spread != 1) {
    stop(paste("Response variable", attr(formula, "response"), 
               "must be a factor with exactly 2 levels."))
  }
  respVec <- respVec - min.respVec
  
  freq.choices <- colSums(matrix(respVec, nrow = N, ncol = K))/N
  loFreq <- min(freq.choices)
  loChoice <- choice.set[which(loFreq == freq.choices)]
  names(freq.choices) <- choice.set
  if (loFreq < 1e-7) {
    cat("Frequencies of alternatives in input data:\n")
    print(prop.table(freq.choices), digits = 4)
    stop(paste("Frequency, in response, of choice:", loChoice, "< 1e-7."))
  }
  
  # Form design matrices 
  formDesignMat <- function(varVec = NULL, includeIntercept = TRUE)
  {
    if (is.null(varVec) && !includeIntercept) return(NULL) 
    fm <- paste(attr(formula, "response"), "~")
    if (!is.null(varVec))
      fm <- paste(fm, paste(varVec, collapse = "+"))
    if (!includeIntercept) fm <- paste(fm, "-1 ")
    else fm <- paste(fm, "+1 ")
    modMat <- model.matrix(as.formula(fm), data)
  } 
  X <- formDesignMat(varVec = attr(formula, "indSpVar"), 
                     includeIntercept = attr(formula, "Intercept"))
  Y <- formDesignMat(varVec = attr(formula, "csvChCoeff"), 
                     includeIntercept = FALSE)
  Z <- formDesignMat(varVec = attr(formula, "csvGenCoeff"), 
                     includeIntercept = FALSE)
  # Determine problem parameters
  response <- respVec
  K <- length(choice.set) # number of choices
  N <- length(response)/(K)       # num of observations
  p <- ifelse(is.null(X), 0, ncol(X)) # num of ind sp variables
  f <- ifelse(is.null(Y), 0, ncol(Y)) # num of ind sp variables
  d <- ifelse(is.null(Z), 0, ncol(Z)) # num of ind sp variables
  nparams <- (K - 1)*p + K*f + d # total number of coeffs
  size <- structure(list(
    "N" = N,
    "K" = K, 
    "p" = p, 
    "f" = f, 
    "d" = d, 
    "nparams" = nparams),
    class = "model.size")
  coeffVec <- mnlogitobject$coefficients #ml.m.w$coefficients
  X <- if (!is.null(X)) X[1:N, , drop=FALSE]   # Matrix of ind sp vars
  # Drop rows for base alternative
  respVec <- respVec[(N + 1):(K*N)]
  response <- respVec
  # Initialize utility matrix: dim(U) = N x K-1
  probMat <- matrix(rep(0, size$N * (size$K-1)), 
                    nrow = size$N, ncol = size$K-1)
  
  # First compute the utility matrix (stored in probMat)
  if (size$p) {
    probMat <- probMat + X %*% matrix(coeffVec[1:((size$K-1) *size$p)],
                                      nrow = size$p, ncol = (size$K-1), byrow=FALSE)
  }
  if (size$f) {
    findYutil<- function(ch_k)
    {
      offset <- (size$K - 1)*size$p
      init <- (ch_k - 1)*size$N + 1
      fin <- ch_k * size$N
      Y[init:fin, , drop=FALSE] %*%
        coeffVec[((ch_k-1)*size$f + 1 + offset):(ch_k*size$f+offset)]
    }
    vec <- as.vector(sapply(c(1:size$K), findYutil))
    vec <- vec - vec[1:size$N]
    
    probMat <- probMat + matrix(vec[(size$N+1):(size$N*size$K)], 
                                nrow = size$N, ncol = (size$K-1), byrow=FALSE)
  }
  if (size$d) {
    probMat <- probMat +
      matrix(Z %*% coeffVec[(size$nparams - size$d + 1):size$nparams],
             nrow = size$N, ncol=(size$K-1), byrow=FALSE)
  }
  # Convert utility to probabilities - use logit formula
  probMat <- exp(probMat)                           # exp(utility)
  baseProbVec <- 1/(1 + rowSums(probMat))           # P_i0
  probMat <- probMat * matrix(rep(baseProbVec, size$K-1),
                              nrow = size$N, ncol = size$K-1) # P_ik
  hess<-1
  # Gradient calculation
  gradient <- if (hess) {
    responseMat <- matrix(response, nrow=size$N, ncol=(size$K-1))
    baseResp <- rep(1, size$N) - rowSums(responseMat)
    
    xgrad <- if (!is.null(X)) {
      as.vector(crossprod(X, responseMat - probMat))
    } else NULL
    xgradmat <- if (!is.null(X)) {
      findXgradmat <- function(ch_p)
      {
        as.vector(X[,ch_p]*(responseMat - probMat))
      }
      inter <- sapply(c(1:size$p), findXgradmat)
      findXgradinter <- function(ch_k)
      {
        matrix(inter[c(((ch_k-1)*size$N+1):(ch_k*size$N)),],nrow=size$N)
      }
      matrix(sapply(c(1:(size$K-1)), findXgradinter),nrow=size$N)
    } else NULL
    ygrad <- if (!is.null(Y)) {
      yresp <- c((baseResp - baseProbVec), (response-as.vector(probMat)))
      findYgrad <- function(ch_k)
      {
        init <- (ch_k - 1)*size$N + 1
        fin <- ch_k * size$N
        crossprod(Y[init:fin, , drop=FALSE], yresp[init:fin])
      }
      as.vector(sapply(c(1:size$K), findYgrad))
    } else NULL
    
    ygradmat <- if (!is.null(Y)) {
      yresp <- c((baseResp - baseProbVec), (response-as.vector(probMat)))
      findYgradmat <- function(ch_k)
      {
        init <- (ch_k - 1)*size$N + 1
        fin <- ch_k * size$N
        inita <- ((ch_k %% size$K==0)*size$K + ch_k %% size$K - 1)*size$N + 1 ##changed in v2 to accomodate interaction terms 
        fina <- ((ch_k %% size$K==0)*size$K + ch_k %% size$K) * size$N ##changed in v2 to accomodate interaction terms 
        Y[c(init:fin)]*yresp[c(inita:fina)] ##changed in v2 to accomodate interaction terms 
      }
      sapply(c(1:(size$K*size$f)), findYgradmat) ##changed in v2 to accomodate interaction terms 
    } else NULL
    if (!is.null(Z)) {
      zgrad <- as.vector(Z) * as.vector(responseMat - probMat)
      zgrad <- colSums(matrix(zgrad, nrow=nrow(Z), ncol=size$d)) #original mlogit
      zgradmat <- matrix(zgrad, nrow=nrow(Z), ncol=size$d) #mlogitjg 
    } else {
      zgrad <- NULL
      zgradmat <- NULL
    }
    gradientmat <- cbind(xgradmat, ygradmat, zgradmat)
    colnames(gradientmat) <- names(mnlogitobject$gradient) #names(ml.m.w$gradient)
 #   c(xgrad, ygrad, zgrad)
    gradientmat
  }
}